import pandas as pd
import numpy as np
import sklearn
from sklearn.model_selection import train_test_split, KFold
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsClassifier
import joblib
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import matplotlib.pylab as pyb


class KELM:
    def __init__(self, kernel='rbf', sigma=1.0, C=88):
        self.kernel = kernel
        self.sigma = sigma
        self.C = C  # 正则化参数
        self.alpha = None
        self.train_X = None
        self.classes_ = None
        # self.encoder = OneHotEncoder(sparse=False)

    # def _rbf_kernel(X, Y, gamma=0.5):
    #     diff = X[:, np.newaxis] - Y[np.newaxis, :]
    #     return np.exp(-gamma * np.sum(diff ** 2, axis=-1))

    def _rbf_kernel(self, X, Y):
        """优化后的RBF核计算"""
        # X_norm = np.sum(X ** 2, axis=1, keepdims=True)
        X_norm = np.sum(X ** 2, axis=1).reshape(-1, 1)
        Y_norm = np.sum(Y ** 2, axis=1)
        # Y_norm = np.sum(Y ** 2, axis=1)
        distance = X_norm + Y_norm - 2 * np.dot(X, Y.T)
        # return np.exp(-distance / (2 * self.sigma ** 2))
        return np.exp(-distance / 1024)

    def _linear_kernel(self, X, Y):
        """线性核函数"""
        return np.dot(X, Y.T)

    def _kernel_matrix(self, X, Y):
        """支持多种核函数"""
        if self.kernel == 'rbf':
            return self._rbf_kernel(X, Y)
        elif self.kernel == 'linear':
            return self._linear_kernel(X, Y)
        else:
            raise ValueError("Unsupported kernel type")

    def fit(self, X, y):
        self.train_X = X
        self.classes_ = np.unique(y)

        # 计算核矩阵
        K = self._kernel_matrix(X, X)

        # 添加正则化项
        I_over_C = np.eye(K.shape[0]) / self.C
        Omega = K + I_over_C

        # 计算权重
        # self.alpha = np.linalg.pinv(Omega) @ y_encoded
        self.alpha = np.linalg.solve(Omega, y)

    def predict(self, X):
        K = self._kernel_matrix(X, self.train_X)
        pred = np.dot(K, self.alpha)
        return (pred > 0.5).astype(int).flatten()
        # pred = K @ self.alpha
        #
        # if len(self.classes_) > 2:
        #     return np.argmax(pred, axis=1)
        # else:
        #     return (pred > 0.5).astype(int).flatten()


class KELMClassifier:
    def __init__(self, path, selectFeatures, selectTargets, resultPath, isRefresh=False, trainRate=0.7):
        self.path = path
        self.selectFeatures = selectFeatures
        self.selectTargets = selectTargets
        self.resultPath = resultPath
        self.isRefresh = isRefresh
        self.trainRate = trainRate
        self.model = None
        self.scaler = StandardScaler()

    def load_data(self):
        # 读取数据文件
        data = pd.read_csv(self.path)
        # 筛选特征和目标
        X = data.iloc[:, self.selectFeatures].values
        y = data.iloc[:, self.selectTargets].values

        y = y.ravel()

        return X, y

    def shuffle_data(self, X, y):
        if self.isRefresh:
            # 如果需要洗牌，打乱数据
            print("Shuffling data...")
            X, y = sklearn.utils.shuffle(X, y)
        return X, y

    def split_data(self, X, y):
        # 按照 trainRate 拆分数据为训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=self.trainRate, random_state=22)

        print(X_train[0])

        return X_train, X_test, y_train, y_test

    def train(self, X_train, y_train):
        # 对数据进行标准化
        X_train_scaled = self.scaler.fit_transform(X_train)
        # X_train_scaled = X_train

        self.model = KELM()
        self.model.fit(X_train_scaled, y_train)

        y_pred = self.model.predict(X_train_scaled)

        # 计算准确率
        accuracy = accuracy_score(y_train, y_pred)
        precision = precision_score(y_train, y_pred, average='weighted')
        recall = recall_score(y_train, y_pred, average='weighted')
        f1 = f1_score(y_train, y_pred, average='weighted')
        return {
            "accuracy": format(accuracy, '.3f'),
            "precision": format(precision, '.3f'),
            "recall": format(recall, '.3f'),
            "f1Score": format(f1, '.3f')
        }

    def evaluate(self, X_test, y_test):
        # 对测试数据进行标准化
        X_test_scaled = self.scaler.transform(X_test)
        # X_test_scaled = X_test
        # 进行预测
        y_pred = self.model.predict(X_test_scaled)
        print(y_pred)

        # 计算准确率
        accuracy = accuracy_score(y_test, y_pred)
        precision = precision_score(y_test, y_pred, average='weighted')
        recall = recall_score(y_test, y_pred, average='weighted')
        f1 = f1_score(y_test, y_pred, average='weighted')
        # 打印评估结果
        print(f"Accuracy: {accuracy:.4f}")
        print(f"Precision: {precision:.4f}")
        print(f"Recall: {recall:.4f}")
        print(f"F1-Score: {f1:.4f}")
        return {
            "accuracy": format(accuracy, '.3f'),
            "precision": format(precision, '.3f'),
            "recall": format(recall, '.3f'),
            "f1Score": format(f1, '.3f')
        }

    def save_model(self):
        # 保存训练好的模型和数据预处理器（Scaler）
        joblib.dump(self.model, f"{self.resultPath}/model.joblib")
        joblib.dump(self.scaler, f"{self.resultPath}/scaler.joblib")

    def cross_validate(self):
        # 加载数据
        X, y = self.load_data()
        # 洗牌数据（如果需要）
        X, y = self.shuffle_data(X, y)

        # 使用 KFold 进行十折交叉验证
        kf = KFold(n_splits=10, shuffle=True)

        fold_results = []

        for train_index, test_index in kf.split(X):
            X_train, X_test = X[train_index], X[test_index]
            y_train, y_test = y[train_index], y[test_index]

            # 训练模型
            train_result = self.train(X_train, y_train)
            # 评估模型
            test_result = self.evaluate(X_test, y_test)
            fold_results.append({
                "trainSet": train_result,
                "testSet": test_result
            })

        # 保存模型
        self.save_model()

        # 计算每个指标的平均值
        avg_results = self.calculate_average(fold_results)

        return avg_results

    def calculate_average(self, fold_results):
        avg_results = {
            "trainSet":{
                "accuracy": 0,
                "precision": 0,
                "recall": 0,
                "f1Score": 0
            },
            "testSet":{
                "accuracy": 0,
                "precision": 0,
                "recall": 0,
                "f1Score": 0
            }
        }

        num_folds = len(fold_results)

        # 对每个指标计算平均值
        for result in fold_results:
            avg_results["testSet"]["accuracy"] += float(result["testSet"]["accuracy"])
            avg_results["testSet"]["precision"] += float(result["testSet"]["precision"])
            avg_results["testSet"]["recall"] += float(result["testSet"]["recall"])
            avg_results["testSet"]["f1Score"] += float(result["testSet"]["f1Score"])

            avg_results["trainSet"]["accuracy"] += float(result["trainSet"]["accuracy"])
            avg_results["trainSet"]["precision"] += float(result["trainSet"]["precision"])
            avg_results["trainSet"]["recall"] += float(result["trainSet"]["recall"])
            avg_results["trainSet"]["f1Score"] += float(result["trainSet"]["f1Score"])

        # 计算平均值
        avg_results["testSet"]["accuracy"] = f"{avg_results['testSet']['accuracy'] / num_folds:.3f}"
        avg_results["testSet"]["precision"] = f"{avg_results['testSet']['precision'] / num_folds:.3f}"
        avg_results["testSet"]["recall"] = f"{avg_results['testSet']['recall'] / num_folds:.3f}"
        avg_results["testSet"]["f1Score"] = f"{avg_results['testSet']['f1Score'] / num_folds:.3f}"

        avg_results["trainSet"]["accuracy"] = f"{avg_results['trainSet']['accuracy'] / num_folds:.3f}"
        avg_results["trainSet"]["precision"] = f"{avg_results['trainSet']['precision'] / num_folds:.3f}"
        avg_results["trainSet"]["recall"] = f"{avg_results['trainSet']['recall'] / num_folds:.3f}"
        avg_results["trainSet"]["f1Score"] = f"{avg_results['trainSet']['f1Score'] / num_folds:.3f}"


        return avg_results

    def run(self):
        # 进行十折交叉验证
        avg_results = self.cross_validate()
        return avg_results


    # def run(self):
    #     # 加载数据
    #     X, y = self.load_data()
    #     # 洗牌数据（如果需要）
    #     X, y = self.shuffle_data(X, y)
    #     # 拆分数据集
    #     X_train, X_test, y_train, y_test = self.split_data(X, y)
    #     # 训练模型
    #     train = self.train(X_train, y_train)
    #     # 评估模型
    #     result = self.evaluate(X_test, y_test)
    #
    #     # 保存模型
    #     self.save_model()
    #
    #     return {
    #         "trainSet": train,
    #         "testSet": result
    #     }


if __name__ == '__main__':
    # 使用示例
    # path = "D:\\Code\\WebProjects\\Dataset\\yds-dataset-analyse\\pythonProject\\test\\2025-03-29_nrurtfb0uim9fuu7.csv"
    path = "D:\\Code\\WebProjects\\Dataset\\yds-dataset-analyse\\pythonProject\\test\\heart_failure_clinical_records_dataset.csv"
    selectFeatures = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
    # selectFeatures = [0, 1, 2, 3, 4, 5,6,7,8,9,10,11,12,13,14,15,16,17,18,19, 20, 21]
    # selectFeatures = [0,1,2, 10,11, 14,15]
    # selectTargets = [22]
    selectTargets = [12]

    resultPath = "D:\\Code\\WebProjects\\Dataset\\yds-dataset-analyse\\pythonProject\\test\\result"
    isRefresh = True
    trainRate = 0.7

    knn_classifier = KELMClassifier(path, selectFeatures, selectTargets, resultPath, isRefresh, trainRate)
    result = knn_classifier.run()
    print(result)
